import os
from os.path import relpath


def find_video_files(directory):
    # 定义常见的视频文件扩展名
    video_extensions = ['.mp4', '.webm']

    video_files = {}

    # 遍历指定目录及其子目录
    for root, dirs, files in os.walk(directory):
        # 过滤出所有视频文件
        videos_in_dir: list[str] = [f for f in files if os.path.splitext(f)[1].lower() in video_extensions]

        if videos_in_dir:
            rel_path = relpath(root, directory)
            # 将每个目录下的视频文件保存在字典中
            video_files[rel_path] = videos_in_dir

    return video_files


def count_videos_count(d):
    length = 0

    for key, value in d.items():
        if not isinstance(value, list):
            length += 1
        else:
            length += len(value)
    return length


# 调用函数，查找指定目录下的视频文件
directory = 'C:\\Users\\chase\\Documents\\Projects\\python\\video\\aesTest'
result = find_video_files(directory)
size = count_videos_count(result)

print(result)
print(f"result size: {size}")
